﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Text;
using RayDen.Library.Components.Surface.Microfacet;
using RayDen.Library.Core;
using RayDen.Library.Core.Components;
using RayDen.Library.Core.Primitives;
using RayDen.Library.Entity.Scene;

namespace RayDen.Library.Components.Surface
{
    public static class FresnelController
    {

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public static RgbSpectrum FrDiel(float coso, float cost, RgbSpectrum etao, RgbSpectrum etat)
        {
            RgbSpectrum Rparl = ((etat * coso) - (etao * cost)) / ((etat * coso) + (etao * cost));
            RgbSpectrum Rperp = ((etao * coso) - (etat * cost)) / ((etao * coso) + (etat * cost));

            return (Rparl * Rparl + Rperp * Rperp) / 2f;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public static RgbSpectrum FrCond(float coso, RgbSpectrum eta, RgbSpectrum k)
        {
            RgbSpectrum tmp = (eta * eta + k * k) * coso * coso;
            RgbSpectrum Rparl2 = (tmp - (2f * eta * coso) + 1) /
            (tmp + (2f * eta * coso) + 1);
            RgbSpectrum tmp_f = eta * eta + k * k;
            RgbSpectrum Rperp2 = (tmp_f - (2f * eta * coso) + coso * coso) / (tmp_f + (2f * eta * coso) + coso * coso);
            return (Rparl2 + Rperp2) / 2f;
        }


        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public static RgbSpectrum Schlick_FresnelTerm(RgbSpectrum Rs, float viewAndHalf)
        {
            RgbSpectrum result = Rs + (RgbSpectrum.UnitSpectrum() - Rs) * (1f - (float)Math.Pow(viewAndHalf, 5f));

            return result;
        }

        public static RgbSpectrum FresnelApproxEta(RgbSpectrum i)
        {
            RgbSpectrum intensity = new RgbSpectrum(i);
            intensity.Clip(0f, 0.999f);
            return (RgbSpectrum.UnitSpectrum() + intensity.Sqrt()) / (RgbSpectrum.UnitSpectrum() - intensity.Sqrt());
        }

        public static RgbSpectrum FresnelApproxK(RgbSpectrum i)
        {
            RgbSpectrum intensity = new RgbSpectrum(i);
            intensity.Clip(0f, 0.999f);
            return 2f * (intensity / (RgbSpectrum.UnitSpectrum() - intensity)).Sqrt();
        }

    }

    public interface IFresnelTerm
    {
        RgbSpectrum Evaluate(float cosi);
    }

    public class FresnelDielectric : IFresnelTerm
    {
        //RgbSpectrum eta = new RgbSpectrum(), k = new RgbSpectrum();

        float eta_o, eta_t;

        public FresnelDielectric(float eo, float et)
        {
            eta_o = eo;
            eta_t = et;
        }

        #region IFresnelTerm Members

        public RgbSpectrum Evaluate(float cosi)
        {
            float coso = MathLab.Clamp(cosi, -1f, 1f);
            bool entering = coso > 0f;
            float eo = eta_o, et = eta_t;
            if (!entering)
                MathLab.Swap(ref eo, ref et);

            float sint = eo / et * MathLab.Sqrt(Math.Max(0f, 1f - coso * coso)); //Snell's Law

            if (sint > 1f)
            {
                return RgbSpectrum.UnitSpectrum();
            }
            else
            {
                float cost = MathLab.Sqrt(Math.Max(0f, 1f - sint * sint));
                return FresnelController.FrDiel(Math.Abs(coso), cost, new RgbSpectrum(eo), new RgbSpectrum(et));
            }
        }

        #endregion
    }


    public class FresnelConductor : IFresnelTerm
    {
        RgbSpectrum eta, k;

        public FresnelConductor(RgbSpectrum e, RgbSpectrum kk)
        {
            eta = new RgbSpectrum(e);
            k = new RgbSpectrum(kk);
        }

        #region IFresnelTerm Members

        public RgbSpectrum Evaluate(float cosi)
        {
            return FresnelController.FrCond(Math.Abs(cosi), eta, k);
        }

        #endregion
    }

    public class FresnelNoOP : IFresnelTerm
    {
        #region IFresnelTerm Members

        public RgbSpectrum Evaluate(float cosi)
        {
            return RgbSpectrum.UnitSpectrum();
        }

        #endregion
    }
    [Flags]
    public enum BxDFTypes : byte
    {
        BSDF_REFLECTION = 1 << 0,
        BSDF_TRANSMISSION = 1 << 1,
        BSDF_DIFFUSE = 1 << 2,
        BSDF_GLOSSY = 1 << 3,
        BSDF_SPECULAR = 1 << 4,
        BSDF_ALL_TYPES = BSDF_DIFFUSE |
                                BSDF_GLOSSY |
                                BSDF_SPECULAR,
        BSDF_ALL_REFLECTION = BSDF_REFLECTION |
                                BSDF_ALL_TYPES,
        BSDF_ALL_TRANSMISSION = BSDF_TRANSMISSION |
                                BSDF_ALL_TYPES,
        BSDF_ALL = BSDF_ALL_REFLECTION |
                                BSDF_ALL_TRANSMISSION
    }

    public abstract class BxDF
    {
        protected BxDFTypes type;

        public BxDFTypes Type
        {
            get
            {
                return type;
            }
        }

        public BxDF(BxDFTypes t)
        {
            type = t;
        }


        public static bool MatchesFlags(BxDF obj, BxDFTypes flags)
        {
            return (flags & obj.type).Equals(flags) || (flags & obj.type).Equals(obj.type);
        }

        public abstract RgbSpectrum F(ref Vector wo, ref  Vector wi);
        public virtual RgbSpectrum Sample(ref Vector wo, out Vector wi, float u1, float u2, out float pdf, out bool specularBounce)
        {
            wi = MC.CosineSampleHemisphere(u1, u2);//TODO REFLECTION COORD SYSTEM
            if (wo.z < 0f) wi.z *= -1f;
            pdf = Pdf(ref wo, ref wi);
            specularBounce = false;
            return F(ref wo, ref wi);
        }


        public virtual float Pdf(ref Vector wo, ref Vector wi)
        {
            return BrdfBase.SameHemisphere(ref wo, ref  wi) ? (float)Math.Abs(wi.z) * MathLab.INVPI : 0f;
        }

        public virtual RgbSpectrum Rho()
        {
            throw new NotImplementedException();
        }

        public virtual RgbSpectrum Rho(Vector wo, int nSamples = 16, float[] samples = null)
        {
            //        if (samples == null) {
            //    samples =
            //        (float *)alloca(2 * nSamples * sizeof(float));
            //    LatinHypercube(samples, nSamples, 2);
            //}
            //RgbSpectrum r = 0.;
            //for (int i = 0; i < nSamples; ++i) {
            //    // Estimate one term of $\rho_{dh}$
            //    Vector wi;
            //    float pdf = 0.f;
            //    RgbSpectrum f =
            //        Sample_f(w, &wi, samples[2*i], samples[2*i+1], &pdf);
            //    if (pdf > 0.) r += f * Math.Abs(wi.z) / pdf;
            //}
            //return r / nSamples; 
            throw new NotImplementedException();
        }


        internal Vector ShadingNormal()
        {
            throw new NotImplementedException();
        }



        internal RgbSpectrum Sample_f(Vector wo, ref Vector wi, float bs1, float bs2, float bcs, out float bsdfPdf, BxDFTypes flags)
        {
            throw new NotImplementedException();
        }
    }

    public class Lambertian : BxDF
    {
        RgbSpectrum R, RPi;

        public Lambertian(RgbSpectrum reflectance)
            : base(
            BxDFTypes.BSDF_REFLECTION | BxDFTypes.BSDF_DIFFUSE)
        {
            R = new RgbSpectrum(reflectance);
            RPi = reflectance * MathLab.INVPI;
        }

        public override RgbSpectrum F(ref Vector wo, ref Vector wi)
        {
            return RPi;
        }

        public override RgbSpectrum Rho(Vector wo, int nSamples = 16, float[] samples = null)
        {
            return R;
        }

        public override RgbSpectrum Rho()
        {
            return R;
        }
    }



    public class OrenNayar : BxDF
    {
        RgbSpectrum R;
        float A, B;

        public OrenNayar(ref RgbSpectrum reflectance, float sig)
            : base(
            BxDFTypes.BSDF_REFLECTION |
            BxDFTypes.BSDF_DIFFUSE)
        {
            R = reflectance;
            float sigma = MathLab.Radians(sig);
            float sigma2 = sigma * sigma;
            A = 1f - (sigma2 / (2f * (sigma2 + 0.33f)));
            B = 0.45f * sigma2 / (sigma2 + 0.09f);
        }

        public override RgbSpectrum F(ref Vector wo, ref Vector wi)
        {
            float sinthetai = BrdfBase.SinTheta(ref wi);
            float sinthetao = BrdfBase.SinTheta(ref wo);
            // Compute cosine term of Oren--Nayar model
            float sinphii = BrdfBase.SinPhi(ref wi), cosphii = BrdfBase.CosPhi(ref wi);
            float sinphio = BrdfBase.SinPhi(ref wo), cosphio = BrdfBase.CosPhi(ref wo);
            float dcos = cosphii * cosphio + sinphii * sinphio;
            float maxcos = Math.Max(0f, dcos);
            // Compute sine and tangent terms of Oren--Nayar model
            float sinalpha, tanbeta;
            if (Math.Abs(BrdfBase.CosTheta(ref wi)) > Math.Abs(BrdfBase.CosTheta(ref wo)))
            {
                sinalpha = sinthetao;
                tanbeta = sinthetai / Math.Abs(BrdfBase.CosTheta(ref wi));
            }
            else
            {
                sinalpha = sinthetai;
                tanbeta = sinthetao / Math.Abs(BrdfBase.CosTheta(ref wo));
            }
            Assert.IsNotNaN(A);
            Assert.IsNotNaN(B);
            Assert.IsNotNaN(maxcos);
            Assert.IsNotNaN(sinalpha);
            Assert.IsNotNaN(tanbeta);

            return R * MathLab.INVPI *
                   (A + B * maxcos * sinalpha * tanbeta);
        }


    }

    public class SpecularTransmission : BxDF
    {
        RgbSpectrum T;
        float etai, etat;
        FresnelDielectric fresnel;

        public SpecularTransmission(RgbSpectrum t, float eta_i, float eta_t)
            : base(BxDFTypes.BSDF_TRANSMISSION | BxDFTypes.BSDF_SPECULAR)
        {
            etai = eta_i;
            etat = eta_t;
            fresnel = new FresnelDielectric(eta_i, eta_t);
            T = t;
        }
        public override RgbSpectrum F(ref Vector wo, ref Vector wi)
        {
            bool entering = BrdfBase.CosTheta(ref wo) > 0f;
            float ei = etai, et = etat;
            if (!entering)
                MathLab.Swap(ref ei, ref  et);
            // Compute transmitted ray direction
            float sini2 = BrdfBase.SinTheta2(ref wo);
            float eta = ei / et;
            float sint2 = eta * eta * sini2;
            // Handle total internal reflection for transmission
            if (sint2 > 1f) return RgbSpectrum.ZeroSpectrum();
            float cost = MathLab.Sqrt(Math.Max(0f, 1f - sint2));
            if (entering) cost = -cost;
            float sintOverSini = eta;
            wi = new Vector(sintOverSini * -wo.x,
                         sintOverSini * -wo.y,
                         cost);
            //var pdf = 1f;
            RgbSpectrum F = fresnel.Evaluate(BrdfBase.CosTheta(ref wo));
            return (ei * ei) / (et * et) * (RgbSpectrum.UnitSpectrum() - F) * T /
                Math.Abs(BrdfBase.CosTheta(ref wi));
        }
    }

    public class SpecularReflection : BxDF
    {
        RgbSpectrum R;
        IFresnelTerm fresnel;

        public SpecularReflection(IFresnelTerm fr, ref RgbSpectrum r)
            : base(BxDFTypes.BSDF_REFLECTION | BxDFTypes.BSDF_SPECULAR)
        {
            this.fresnel = fr;
            this.R = new RgbSpectrum(r);
        }

        public override float Pdf(ref Vector wo, ref Vector wi)
        {
            return 0f;
        }

        public override RgbSpectrum F(ref Vector wo, ref Vector wi)
        {
            return RgbSpectrum.ZeroSpectrum();
        }

        public override RgbSpectrum Sample(ref Vector wo, out Vector wi, float u1, float u2, out float pdf, out bool specularBounce)
        {
            wi = new Vector(-wo.x, -wo.y, wo.z);
            pdf = 1f;
            specularBounce = true;
            return fresnel.Evaluate(BrdfBase.CosTheta(ref wo)) * R / Math.Abs(BrdfBase.CosTheta(ref wi));
        }
    }

    public class MicrofacetBxDF : BxDF
    {
        RgbSpectrum R;
        MicrofacetDistribution distribution;
        IFresnelTerm fresnel;

        public MicrofacetBxDF(ref RgbSpectrum reflectance,
                       IFresnelTerm f,
                       MicrofacetDistribution d)
            : base(BxDFTypes.BSDF_REFLECTION | BxDFTypes.BSDF_GLOSSY)
        {
            R = reflectance;
            distribution = d;
            fresnel = f;
        }

        public float G(ref Vector wo, ref Vector wi, ref Vector wh)
        {
            float NdotWh = Math.Abs(BrdfBase.CosTheta(ref wh));
            float NdotWo = Math.Abs(BrdfBase.CosTheta(ref wo));
            float NdotWi = Math.Abs(BrdfBase.CosTheta(ref wi));
            float WOdotWh = MathLab.AbsDot(ref wo, ref wh);
            return Math.Min(1f, Math.Min((2f * NdotWh * NdotWo / WOdotWh),
                            (2f * NdotWh * NdotWi / WOdotWh)));
        }

        public override RgbSpectrum F(ref Vector wo, ref Vector wi)
        {
            float cosThetaO = Math.Abs(BrdfBase.CosTheta(ref wo));
            float cosThetaI = Math.Abs(BrdfBase.CosTheta(ref wi));
            Vector wh = (wi + wo).Normalize();
            float cosThetaH = (wi & wh);
            RgbSpectrum F = fresnel.Evaluate(cosThetaH);
            return R * distribution.D(ref wh) * G(ref wo, ref wi, ref wh) * F /
                 (4f * cosThetaI * cosThetaO);
        }

        public override float Pdf(ref Vector wo, ref Vector wi)
        {
            if (!BrdfBase.SameHemisphere(ref wo, ref wi)) return 0f;
            return distribution.Pdf(ref wo, ref wi);
        }

        public override RgbSpectrum Sample(ref Vector wo, out Vector wi, float u1, float u2, out float pdf, out bool specularBounce)
        {
            distribution.Sample_f(ref wo, out wi, u1, u2, out pdf);
            specularBounce = false;
            if (!BrdfBase.SameHemisphere(ref wo, ref wi)) return new RgbSpectrum(0f);
            return F(ref wo, ref wi);
        }
    }



    public class Bsdf
    {
        public Vector WorldToLocal(ref Vector v)
        {
            return new Vector(Vector.Dot(ref v, ref  sn), Vector.Dot(ref v, ref tn), Vector.Dot(ref v, ref nn));
        }
        public Vector LocalToWorld(Vector v)
        {
            return new Vector(sn.x * v.x + tn.x * v.y + nn.x * v.z,
                               sn.y * v.x + tn.y * v.y + nn.y * v.z,
                               sn.z * v.x + tn.z * v.y + nn.z * v.z);
        }
        protected List<BxDF> bxdfs;
        protected internal Vector tn, sn, ng;


        public Vector nn;
        public Bsdf(ref Normal geoN, ref Normal shadeN)
        {
            nn = shadeN.ToVec();
            nn.Normalize();
            ng = geoN.ToVec();
            ng.Normalize();

            Vector.CoordinateSystem(ref nn, out sn, out tn);
            bxdfs = new List<BxDF>();
        }

        public int NumComponents(BxDFTypes flags)
        {
            var res = bxdfs.Count(p => BxDF.MatchesFlags(p, flags));
            return res;
        }

        public RgbSpectrum Sample_f(ref Vector woW, out Vector wiW,
                                    float u1, float u2, float u3, out float pdf, out BxDFTypes? sampled,
                                    BxDFTypes flags = BxDFTypes.BSDF_ALL,
                                    BxDFTypes? sampledType = null)
        {
            BxDFTypes fl = flags;
            int matchingComps = NumComponents(fl);
            if (matchingComps == 0)
            {
                pdf = 0f;
                wiW = Vector.Zero;
                sampled = null;
                return new RgbSpectrum(0f);
            }
            int which = (int)Math.Min((u3 * matchingComps),
                matchingComps - 1);
            BxDF bxdf = null;
            int count = which;
            for (int i = 0; i < bxdfs.Count; ++i)
                if (BxDF.MatchesFlags(bxdfs[i], fl))
                    if (count-- == 0)
                    {
                        bxdf = bxdfs[i];
                        break;
                    }
            Vector wi = new Vector();
            Vector wo = WorldToLocal(ref woW);
            wiW = Vector.Zero;
            pdf = 0f;
            bool specularBounce;
            RgbSpectrum f = bxdf.Sample(ref wo, out wi, u1, u2, out pdf, out specularBounce);
            sampled = bxdf.Type;
            if (pdf > 0f && pdf < MathLab.Epsilon) return RgbSpectrum.ZeroSpectrum();
            if (sampledType != null) sampledType = bxdf.Type;
            wiW = LocalToWorld(wi);

            if ((!BxDF.MatchesFlags(bxdf, BxDFTypes.BSDF_SPECULAR)) && matchingComps > 1)
            {
                for (int i = 0; i < bxdfs.Count; ++i)
                {
                    if (bxdfs[i] != bxdf &&
                        BxDF.MatchesFlags(bxdfs[i], fl))
                        pdf += bxdfs[i].Pdf(ref wo, ref wi);
                }
            }
            if (matchingComps > 1) pdf /= matchingComps;
            // Compute value of BSDF for sampled direction
            if (BxDF.MatchesFlags(bxdf, BxDFTypes.BSDF_SPECULAR))
            //if ((bxdf.Type & BxDFType.BSDF_SPECULAR) == 0)
            {
                f = RgbSpectrum.ZeroSpectrum();
                if ((wiW & ng) * (woW & ng) > 0f)
                // ignore BTDFs
                {
                    fl = fl & ~BxDFTypes.BSDF_TRANSMISSION;
                }
                else
                    // ignore BRDFs
                    fl = (fl & ~BxDFTypes.BSDF_REFLECTION);
                for (int i = 0; i < bxdfs.Count; ++i)
                    if (BxDF.MatchesFlags(bxdfs[i], fl))
                        f += bxdfs[i].F(ref wo, ref wi);
                f = f / pdf;
            }
            return f;
        }

        public float Pdf(ref Vector woW, ref Vector wiW, BxDFTypes flags)
        {
            if (bxdfs.Count == 0) return 0f;
            Vector wo = WorldToLocal(ref woW);
            var wi = WorldToLocal(ref wiW);
            float pdf = 0f;

            int matchingComps = 0;
            for (int i = 0; i < bxdfs.Count; ++i)
                if (BxDF.MatchesFlags(bxdfs[i], flags))
                {
                    ++matchingComps;
                    pdf += bxdfs[i].Pdf(ref wo, ref wi);
                }
            return matchingComps > 0 ? pdf / matchingComps : 0f;
        }


        public RgbSpectrum f(ref Vector woW, ref Vector wiW, BxDFTypes flags)
        {
            Vector wi = WorldToLocal(ref wiW);
            var wo = WorldToLocal(ref woW);
            if ((wiW & ng) * (woW & ng) > 0)
                // ignore BTDFs
                flags = flags & ~BxDFTypes.BSDF_TRANSMISSION;
            else
                // ignore BRDFs
                flags = flags & ~BxDFTypes.BSDF_REFLECTION;
            RgbSpectrum f = RgbSpectrum.ZeroSpectrum();
            for (int i = 0; i < bxdfs.Count; ++i)
                if (BxDF.MatchesFlags(bxdfs[i], flags))
                    f += bxdfs[i].F(ref wo, ref wi);
            return f;
        }
        public RgbSpectrum rho(BxDFTypes flags)
        {
            RgbSpectrum ret = new RgbSpectrum();
            for (int i = 0; i < bxdfs.Count; ++i)
                if (BxDF.MatchesFlags(bxdfs[i], flags))
                    ret += bxdfs[i].Rho();
            return ret;
        }
        public RgbSpectrum rho(Vector wo, BxDFTypes flags)
        {
            RgbSpectrum ret = new RgbSpectrum();
            for (int i = 0; i < bxdfs.Count; ++i)
                if (BxDF.MatchesFlags(bxdfs[i], flags))
                    ret += bxdfs[i].Rho(wo);
            return ret;
        }

        //public RgbSpectrum Sample_f(ref Vector wo, out Vector wi,
        //                    BxDFType flags = BxDFType.BSDF_ALL,
        //                    BxDFType? sampledType = null)
        //{

        //}

        public void Add(BxDF bxdf)
        {
            this.bxdfs.Add(bxdf);
        }
    }
}
